import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent / "FoundationPose-main"))
from estimater import *
from datareader import *
import argparse
import trimesh
import os
import numpy as np
import cv2
import imageio
import logging
import datetime
import open3d as o3d
import nvdiffrast.torch as dr

class PoseTracker:
    def __init__(self, mesh_file, debug_dir, debug=3, est_refine_iter=5, track_refine_iter=2):
        self.first_call = True
        self.debug = debug
        self.debug_dir = debug_dir
        self.est_refine_iter = est_refine_iter
        self.track_refine_iter = track_refine_iter
        
        mesh = trimesh.load(mesh_file)
        self.to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
        self.bbox = np.stack([-extents/2, extents/2], axis=0).reshape(2,3)
        
        scorer = ScorePredictor()
        refiner = PoseRefinePredictor()
        glctx = dr.RasterizeCudaContext()
        verts = mesh.vertices.astype(np.float32)
        normals = mesh.vertex_normals.astype(np.float32)

        self.est = FoundationPose(model_pts=verts, model_normals=normals, mesh=mesh, scorer=scorer, refiner=refiner, debug_dir=debug_dir, debug=debug, glctx=glctx)
        logging.info("estimator initialization done")
    
    def find_pose(self, rgb, depth, mask, intrin_mat):
        # print("type", intrin_mat.dtype)
        # intrin_mat = torch.as_tensor(intrin_mat, dtype=torch.float64)
        # intrin_mat = intrin_mat.astype(float64)
        pose = self.est.register(K=intrin_mat, rgb=rgb, depth=depth, ob_mask=mask, iteration=self.est_refine_iter)
        if self.debug >= 3:
            m = self.est.mesh.copy()
            m.apply_transform(pose) # maybe not
            m.export(f'{self.debug_dir}/model_tf.obj')
            xyz_map = depth2xyzmap(depth, intrin_mat)
            valid = depth >= 0.001
            pcd = toOpen3dCloud(xyz_map[valid], rgb[valid])
            o3d.io.write_point_cloud(f'{self.debug_dir}/scene_complete.ply', pcd)
        return pose
    
    def track_pose(self, rgb, depth, intrin_mat):
        pose = self.est.track_one(rgb=rgb, depth=depth, K=intrin_mat, iteration=self.track_refine_iter)
        return pose
    
    def get_pose(self, rgb, depth, intrin_mat, extrin_mat, mask=None, frame_id=None):
        assert frame_id is not None, f"frame_id is required, got {frame_id}"
        assert intrin_mat is not None, f"intrin_mat is required , got {intrin_mat}"
        if self.first_call:
            assert mask is not None, f"mask is required for first call, got {mask}"
            pose = self.find_pose(rgb, depth, mask, intrin_mat)
            pose = extrin_mat @ pose
            self.first_call = False
        else:
            pose = self.track_pose(rgb, depth, intrin_mat)
        
        np.savetxt(f'{self.debug_dir}/ob_in_cam/{frame_id}.txt', pose.reshape(4,4))
        
        if self.debug >= 1:
            center_pose = pose @ np.linalg.inv(self.to_origin)
            center_pose = pose
            vis = draw_posed_3d_box(intrin_mat, img=rgb, ob_in_cam=center_pose, bbox=self.bbox)
            vis = draw_xyz_axis(rgb, ob_in_cam=center_pose, scale=0.1, K=intrin_mat, thickness=3, transparency=0, is_input_rgb=True)
            os.makedirs(f'{self.debug_dir}/display', exist_ok=True)
            cv2.imwrite(f'{self.debug_dir}/display/{frame_id}.png', vis[...,::-1])
        
        if self.debug >= 2:
            imageio.imwrite(f'{self.debug_dir}/track_vis/{frame_id}.png', vis)
        
        return pose

def poseTrack(rgb, depth, intrin_mat, extrin_mat, mask, mesh_file, frame_id, debug=3, est_refine_iter=5, track_refine_iter=2):
    assert frame_id is not None, f"frame_id is required, got {frame_id}"
    debug_dir = Path(__file__).parent.parent / "zarrtraj" / "foundationPose"
    if not hasattr(poseTrack,'_tracker'):
        poseTrack._tracker={}
    if  not poseTrack._tracker.get(str(mesh_file)):
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        time_based_dir = debug_dir / timestamp
        (time_based_dir / 'track_vis').mkdir(parents=True, exist_ok=True)
        (time_based_dir / 'ob_in_cam').mkdir(parents=True, exist_ok=True)
        poseTrack._tracker[str(mesh_file)] = PoseTracker(str(mesh_file), str(time_based_dir), debug, est_refine_iter, track_refine_iter)
    pose = poseTrack._tracker[str(mesh_file)].get_pose(rgb, depth, intrin_mat, extrin_mat, mask, frame_id)
    return pose
